#!/usr/bin/env python3
"""
Comprehensive script to check for ALL unclosed environments in LaTeX files.
This should find the root cause of progressive indentation.
"""
import os
import re
import glob

def check_environment_balance(filepath, env_types):
    """Check if all specified environments are properly balanced in a file."""
    with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
        content = f.read()
    
    results = {}
    
    for env in env_types:
        begins = re.findall(rf'\\begin\{{{env}\}}', content)
        ends = re.findall(rf'\\end\{{{env}\}}', content)
        balance = len(begins) - len(ends)
        
        results[env] = {
            'begins': len(begins),
            'ends': len(ends),
            'balance': balance
        }
    
    # Also check for any generic \begin{} without corresponding \end{}
    all_begins = re.findall(r'\\begin\{([^}]+)\}', content)
    all_ends = re.findall(r'\\end\{([^}]+)\}', content)
    
    # Count occurrences of each environment type
    begin_counts = {}
    end_counts = {}
    
    for env in all_begins:
        begin_counts[env] = begin_counts.get(env, 0) + 1
    
    for env in all_ends:
        end_counts[env] = end_counts.get(env, 0) + 1
    
    # Find all environment types that appear in the file
    all_envs = set(all_begins) | set(all_ends)
    
    generic_results = {}
    for env in all_envs:
        begins = begin_counts.get(env, 0)
        ends = end_counts.get(env, 0)
        balance = begins - ends
        if balance != 0:  # Only report imbalanced ones
            generic_results[env] = {
                'begins': begins,
                'ends': ends,
                'balance': balance
            }
    
    return {
        'file': filepath,
        'specific_envs': results,
        'all_imbalanced_envs': generic_results
    }

def main():
    # Common environments that could cause indentation issues
    problematic_envs = [
        'itemize', 'enumerate', 'description',
        'center', 'flushleft', 'flushright',
        'minipage', 'adjustbox', 'quote', 'quotation',
        'algorithm', 'algorithmic', 'align', 'equation',
        'table', 'tabular', 'figure'
    ]
    
    # Check all .tex files
    tex_files = glob.glob('part*/*.tex', recursive=True)
    tex_files.append('gsi_technical_report.tex')  # Include main file
    tex_files.sort()
    
    print("=" * 80)
    print("Comprehensive LaTeX Environment Balance Analysis")
    print("=" * 80)
    
    total_problems = 0
    cumulative_imbalance = {}  # Track cumulative imbalance across all files
    
    for tex_file in tex_files:
        if os.path.exists(tex_file):
            result = check_environment_balance(tex_file, problematic_envs)
            
            file_has_problems = False
            
            # Check specific environments
            for env, data in result['specific_envs'].items():
                if data['balance'] != 0:
                    file_has_problems = True
                    cumulative_imbalance[env] = cumulative_imbalance.get(env, 0) + data['balance']
            
            # Check all imbalanced environments found
            for env, data in result['all_imbalanced_envs'].items():
                file_has_problems = True
                cumulative_imbalance[env] = cumulative_imbalance.get(env, 0) + data['balance']
            
            if file_has_problems:
                print(f"\n⚠️  PROBLEMS in {result['file']}:")
                
                # Print specific environment problems
                for env, data in result['specific_envs'].items():
                    if data['balance'] != 0:
                        print(f"   {env}: begins={data['begins']}, ends={data['ends']}, balance={data['balance']:+d}")
                
                # Print other imbalanced environments
                for env, data in result['all_imbalanced_envs'].items():
                    if env not in problematic_envs:  # Don't double-print
                        print(f"   {env}: begins={data['begins']}, ends={data['ends']}, balance={data['balance']:+d}")
                
                total_problems += 1
            else:
                print(f"✓ {result['file']} - All environments balanced")
    
    print("\n" + "=" * 80)
    print("CUMULATIVE ANALYSIS:")
    print("=" * 80)
    
    if cumulative_imbalance:
        print("Environment imbalances across ALL files:")
        for env, balance in sorted(cumulative_imbalance.items()):
            if balance != 0:
                status = "UNCLOSED" if balance > 0 else "EXTRA CLOSES"
                print(f"   {env}: {balance:+d} ({status})")
                
                if balance > 0:
                    print(f"      → This causes {balance} level(s) of permanent left indentation!")
    else:
        print("✅ No environment imbalances found.")
    
    print(f"\nFiles with problems: {total_problems}/{len(tex_files)}")
    
    # Calculate expected cumulative indentation
    total_indentation_levels = sum(max(0, balance) for balance in cumulative_imbalance.values())
    if total_indentation_levels > 0:
        print(f"\n🔥 CRITICAL: Expected cumulative indentation levels: {total_indentation_levels}")
        print("This explains why later chapters have progressively worse left indentation!")
    
    print("=" * 80)

if __name__ == "__main__":
    main()